﻿using System;
using System.Collections.Generic;
using System.Text;
using System.Reflection;
using System.Collections.ObjectModel;
using System.IO;

namespace UndisposedViewer
{
    class FinalizerTypeFinder : MarshalByRefObject
    {
        Dictionary<string, IEnumerable<Type>> assemblyTypesCache = new Dictionary<string, IEnumerable<Type>>();
        Dictionary<string, object> analyzedAssemblies = new Dictionary<string, object>();

        bool recursiveSearch;
        string basePath = "";

        public FinalizerTypeFinder()
        {
            AppDomain.CurrentDomain.ReflectionOnlyAssemblyResolve += ResolveDependentAssemblies;
        }

        public override object InitializeLifetimeService()
        {
            return null;
        }

        public IEnumerable<string> Find(string assemblyPath, bool recursiveSearch)
        {
            List<string> finalizableTypeNames = new List<string>();

            analyzedAssemblies.Clear();
            exceptionsDuringFind.Clear();
            basePath = System.IO.Path.GetDirectoryName(assemblyPath);

            var assembly = Assembly.ReflectionOnlyLoadFrom(assemblyPath);
            var assemblyName = assembly.GetName();
            this.recursiveSearch = recursiveSearch;

            var finalizableTypes = FindFinalizableTypes(assemblyName);
            List<Type> cache = new List<Type>();

            foreach (Type type in finalizableTypes)
            {
                cache.Add(type);
                //yield return type.FullName;
                finalizableTypeNames.Add(type.FullName);
            }

            AddToCache(assemblyName.FullName, cache);

            if (exceptionsDuringFind.Count > 0)
            {
                throw new Exception("Could not find all finalizable types. " + Environment.NewLine +
                    string.Join(Environment.NewLine, exceptionsDuringFind.ConvertAll(ex => ex.ToString()).ToArray()));
            }
            finalizableTypeNames.Remove("System.Object");
            return finalizableTypeNames;
        }

        List<Exception> exceptionsDuringFind = new List<Exception>();

        private IEnumerable<Type> FindFinalizableTypes(AssemblyName assemblyName)
        {
            IEnumerable<Type> types;
            if (IsPresentInCache(assemblyName.FullName, out types))
            {
                foreach (Type t in types)
                {
                    yield return t;
                }
            }
            else
            {
                Assembly assembly = null;
                try
                {
                    assembly = Assembly.ReflectionOnlyLoad(assemblyName.FullName);
                }
                catch (Exception)
                {
                    try
                    {
                        assembly = Assembly.ReflectionOnlyLoadFrom(
                            System.IO.Path.Combine(basePath, assemblyName.Name + ".dll"));
                    }
                    catch (Exception e)
                    {
                        exceptionsDuringFind.Add(e);
                    }
                }

                if (assembly != null)
                {
                    foreach (Type t in assembly.GetTypes())
                    {
                        if (HasFinalizer(t))
                            yield return t;
                    }

                    if (recursiveSearch)
                    {
                        foreach (AssemblyName a in assembly.GetReferencedAssemblies())
                        {
                            if (!IsAlreadyAnalyzed(a.FullName))
                            {
                                analyzedAssemblies.Add(a.FullName, null);
                                IEnumerable<Type> finalizableTypes = FindFinalizableTypes(a);

                                foreach (Type t in finalizableTypes)
                                    yield return t;

                                AddToCache(a.FullName, new List<Type>(finalizableTypes));
                            }
                        }
                    }
                }
            }
        }

        private bool IsAlreadyAnalyzed(string assemblyName)
        {
            return analyzedAssemblies.ContainsKey(assemblyName);
        }

        private void AddToCache(string assemblyName, IEnumerable<Type> types)
        {
            if (!assemblyTypesCache.ContainsKey(assemblyName))
                assemblyTypesCache.Add(assemblyName, types);
        }

        private bool IsPresentInCache(string assemblyName, out IEnumerable<Type> types)
        {
            return assemblyTypesCache.TryGetValue(assemblyName, out types);
        }

        private bool HasFinalizer(Type type)
        {
            const string finalizerMethodName = "Finalize";
            return type.GetMethod(finalizerMethodName, BindingFlags.NonPublic | BindingFlags.DeclaredOnly | BindingFlags.Instance) != null;
        }

        public Assembly ResolveDependentAssemblies(object sender, ResolveEventArgs args)
        {
            try
            {
                return Assembly.ReflectionOnlyLoad(args.Name);
            }
            catch (Exception e)
            {
                return Assembly.ReflectionOnlyLoadFrom(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, args.Name));
            }
        }
    }
}
